Replace lazy MFSDP context initialization with an explicit scope - #6190
Replace lazy MFSDP context initialization with an explicit scope#6190wujingyue wants to merge 6 commits into
Conversation
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
wujingyue
changed the title
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
| visited: set[FsdpModule] = set() | ||
| root_modules: list[FsdpModule] = [] | ||
| for module in reversed(self._registered_modules): | ||
| if module in visited: | ||
| continue | ||
| root_modules.append(module) | ||
| _collect_fsdp_modules_under(cast(nn.Module, module), visited) | ||
| root_modules.reverse() | ||
| for root_module in root_modules: | ||
| root_module._is_root = True |
There was a problem hiding this comment.
Note to myself: there's an easy way to compute root_modules without assuming the order of registered_modules
|
|
||
|
|
||
| @contextmanager | ||
| def fully_shard_context(device: torch.device | None = None) -> Iterator[FsdpContext]: |
There was a problem hiding this comment.
Example of silent misbehavior:
with fully_shard_context(device=device) as ctx1:
fully_shard(model.inner, mesh=mesh, placements=placements)
with fully_shard_context(device=device) as ctx2:
fully_shard(model, mesh=mesh, placements=placements)
# ctx2.finalize() silently includes model.inner in its forward_order,
# while model.inner._is_root=True from ctx1 still holds.
Maybe worth adding a guard or UT that detects when an FsdpModule being registered already belongs to another context, and raises explicitly.
Also, If we require all fully_shard calls to live under a single continuous fully_shard_context, does that make it harder to align with the FSDP2 API? (I recall the plan was to achieve compatibility via a wrapper)
There was a problem hiding this comment.
Maybe worth adding a guard or UT that detects when an FsdpModule being registered already belongs to another context, and raises explicitly.
Yep! Good idea.
| for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): | ||
| is_last = microbatch_index == num_microbatches - 1 | ||
| with microbatch(model, is_last=is_last): | ||
| with microbatch(context, is_last=is_last): |
There was a problem hiding this comment.
If the context is already attached to the model, would it be cleaner to retrieve it via something like model.get_context()?
There was a problem hiding this comment.
If model is an FsdpModule, it already has model.context.
Otherwise, I’d prefer fewer places to track global state. Since this PR introduces fully_shard_context, we can reuse it for all knobs shared across FsdpModules. Along those lines, perhaps fully_shard_optimizer should take context rather than model as well.
2 participants
What does this PR do?
Why
MFSDP v2 has lazily initialized its runtime context on the first forward to stay close to PyTorch FSDP2. That defers stream creation, root discovery, module naming, and prefetch-order construction until runtime, adding lifecycle checks and coupling context setup to forward hooks.
#6187 is a concrete example of the resulting complexity. Because the reduce-scatter stream is unavailable when
fully_shard()constructs parameter groups,main_gradmust be materialized lazily so its allocation occurs on the stream that later releases it. With an eagerly initialized context, stream-dependent state can instead be constructed against the intended stream up front.What changes
This PR introduces an explicit construction scope:
Every
fully_shard()call must run insidefully_shard_context(). Calling it without an active context raises a hard error. Nested construction scopes are also rejected.The context eagerly creates its communication streams. Enclosed FSDP modules naturally inherit that context and register with it. When the scope exits, the context discovers roots, assigns module names, and builds the static forward and backward prefetch orders. Forward execution before scope exit also raises a hard error.
microbatch()now takes the context directly, so last-microbatch state is context-scoped rather than discovered by walking from a root module.Other benefits
fully_shard_context()participate in one cross-root overlap schedule. Users no longer need an artificialfully_shard(root_model)call merely to make independently sharded children share streams and overlap communication.Issue tracking
Linked issue:
Contribution process
Pre-checks
Validation
/opt/venv/bin/python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2isorton all changed Python files.black --checkon all changed Python files.ruff checkon all changed Python files.python -m compileallon the changed implementation and tests.git diff --check.